package phantm.phases

import phantm._
import phantm.symbols._
import phantm.cfg.{ LabeledDirectedGraphImp, ControlFlowGraph }
import phantm.cfg.EdgeImp
import phantm.cfg.Trees._
import java.io.{ PrintStream, FileOutputStream }
import phantm.types._
import phantm.cfg.ControlFlowGraph
import phantm.util.PhasesUtils

object FunctionCallGraphPhase extends Phase {
  def name = "Generating function call graph"
  def description = "Generating the function call graph starting from main"

  def run(ctx: PhasesContext): PhasesContext = {
    if (!Settings.get.exportPaths.exportFCGPath.isEmpty) {
      val generator = FunctionCallGraphGenerator(ctx)
      generator.execute
      val functionCallGraph = generator.FunctionCallGraph
      functionCallGraph.writeDottyToFile(Settings.get.exportPaths.exportFCGPath.get, "Function Call Graph")
    }
    ctx
  }
}

case class FunctionCallGraphGenerator(ctx: PhasesContext) {

  object FunctionCallGraph extends LabeledDirectedGraphImp[String] {

    type AVertex = FunctionSymbol

    val entry = newVertex
    entry.name = "\"[main]\""

    var symToV = Map[FunctionSymbol, Vertex]()
    var vToSym = Map[Vertex, FunctionSymbol]()
    var usedNames: Set[String] = Set()

    def addNode(functionSymbol: AVertex, className: String): Vertex = {
      symToV.get(functionSymbol) match {
        case Some(v) => v
        case None =>
          val vertex = newVertex
          val name = functionSymbol.name
          if (usedNames contains name) {
            vertex.name = className + "_" + name
          } else {
            vertex.name = name
            usedNames += name
          }
          symToV += (functionSymbol -> vertex)
          vToSym += (vertex -> functionSymbol)
          this += vertex
          vertex
      }
    }

    def addEdge(from: AVertex, to: AVertex, className: String, isStatic: Boolean) = {
      val fromVertex = addNode(from, className)
      val toVertex = addNode(to, className)
      if (isStatic) {
        this += (fromVertex, className, toVertex, "lightblue")
      } else {
        this += (fromVertex, className, toVertex, "yellow")
      }
    }
  }

  def execute = {

    val helper = new PhasesUtils(ctx)
    val mainFunction: FunctionSymbol = new FunctionSymbol("\"[main]\"")
    var functionsMap: Map[FunctionSymbol, Set[(FunctionSymbol, Boolean, String)]] = Map()
    var functionsMapFromMain: Map[FunctionSymbol, Set[FunctionSymbol]] = Map()
    var visited: Set[FunctionSymbol] = Set()

    def computeFunctionCall(aFunction: FunctionSymbol): Unit = {

      val optionCfg: Option[ControlFlowGraph] = {
        if (aFunction.name equals mainFunction.name) helper.getCFG(None) else
          helper.getCFG(Some(aFunction))
      }
      optionCfg match {
        case Some(cfg) =>
          var edges = cfg.E

          def handleMethodCall(optionClassSymbol: Option[ClassSymbol], optionMethod: Option[MethodSymbol], staticFlag: Boolean): Unit = {
            (optionClassSymbol, optionMethod) match {
              case (Some(classSymbol), Some(method)) if (method.userland) =>
                val value: Set[(FunctionSymbol, Boolean, String)] = functionsMap.getOrElse(aFunction, Set())
                val mainValue: Set[FunctionSymbol] = functionsMapFromMain.getOrElse(aFunction, Set())
                functionsMap += aFunction -> (value + Tuple3(method, staticFlag, classSymbol.name))
                functionsMapFromMain += aFunction -> (mainValue + method)
              case _ =>
            }
          }

          while (!edges.isEmpty) {
            val edge = edges.head
            val fromVertex = edge.v1
            val vertexEnvironmentMap = ctx.cfgsWithTypes(cfg)
            val environment = vertexEnvironmentMap(fromVertex)
            edge.lab match {
              case Assign(variable, MethodCall(receiver, methodId, _)) =>
                val optionClassSymbol = helper.classFromSimpleValue(receiver, environment)
                handleMethodCall(optionClassSymbol, helper.lookupMethod(optionClassSymbol, methodId), false)
              case Assign(variable, StaticMethodCall(classRef, methodId, _)) => classRef match {
                case ClassRefFixed(classSymbol) => handleMethodCall(Some(classSymbol), helper.lookupMethod(Some(classSymbol), methodId), true)
                case _ =>
              }
              case Assign(variable, New(classRef, parameters)) => classRef match {
                case ClassRefFixed(classSymbol) => handleMethodCall(Some(classSymbol), helper.lookupConstructor(classSymbol), false)
                case _ =>
              }
              case _ =>
            }
            edges = edges - edge
          }
        case None =>
      }
    }

    computeFunctionCall(mainFunction)
    while (!functionsMap.isEmpty) {
      val fromFunction = functionsMap.head._1
      for (toFunctionPair <- functionsMap.head._2) {
        val toFunction = toFunctionPair._1
        FunctionCallGraph.addEdge(fromFunction, toFunction, toFunctionPair._3, toFunctionPair._2)
        if (!(visited contains toFunction)) {
          computeFunctionCall(toFunction)
          visited += toFunction
        }
      }
      functionsMap = functionsMap - fromFunction
    }
    ctx.results.functionCallFromMain = functionsMapFromMain
  }
}
